import os
import time
import cv2
import matplotlib
from utils import image_util
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt

from ultralytics import YOLO
from torchvision import datasets, transforms
from sklearn.model_selection import train_test_split
from PIL import Image


# 创建 YOLO 格式数据集目录
def create_yolo_dataset_dir(data_dir='./mnist_yolo'):
    """创建 YOLO 格式数据集所需的目录结构"""
    os.makedirs(data_dir, exist_ok=True)
    for split in ['train', 'val', 'test']:
        os.makedirs(f'{data_dir}/{split}/images', exist_ok=True)
        os.makedirs(f'{data_dir}/{split}/labels', exist_ok=True)
    return data_dir


# 转换 MNIST 数据集为 YOLO 格式
def convert_mnist_to_yolo(data_dir='./mnist_yolo'):
    """将 MNIST 数据集转换为 YOLO 格式"""
    # 创建目录
    yolo_dir = create_yolo_dataset_dir(data_dir)

    #  transforms = transforms.Compose([
    #     transforms.Resize((64, 64)),  # 调整为64x64，比原始MNIST大
    #     transforms.Pad(4, padding_mode='constant', fill=0),  # 添加填充
    #     transforms.RandomRotation(15),  # 随机旋转
    #     transforms.RandomCrop(64),  # 随机裁剪
    #     transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),  # 随机平移
    #     transforms.ToTensor(),
    #     transforms.Normalize((0.1307,), (0.3081,))  # MNIST的均值和标准差
    # ])

    # 定义数据转换
    transform = transforms.Compose([
        # transforms.Resize((64, 64)),
        transforms.Pad(padding=50, fill=0),
        transforms.RandomRotation(15),  # 随机旋转
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    # 加载 MNIST 数据集
    train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST('./data', train=False, transform=transform)

    # 划分训练集和验证集
    train_indices, val_indices = train_test_split(
        list(range(len(train_dataset))), test_size=0.2, random_state=42
    )

    # 转换训练集
    for i, (image, label) in enumerate([train_dataset[idx] for idx in train_indices]):
        save_mnist_sample(image, label, f'{yolo_dir}/train', i)

    # 转换验证集
    for i, (image, label) in enumerate([train_dataset[idx] for idx in val_indices]):
        save_mnist_sample(image, label, f'{yolo_dir}/val', i)

    # 转换测试集
    for i, (image, label) in enumerate(test_dataset):
        save_mnist_sample(image, label, f'{yolo_dir}/test', i)

    # 创建数据配置文件
    create_data_yaml(yolo_dir)

    return yolo_dir



def save_mnist_sample(image_tensor, label, split_dir, index):
    image_path = f"{split_dir}/images/{index}.jpg"
    image_util.save_image_tensor(image_tensor, image_path)

    # 创建 YOLO 格式标签文件
    # MNIST 图像中数字居中，占大部分区域
    # 格式: class x_center y_center width height (归一化到 0-1)
    x_center = 0.5
    y_center = 0.5
    width = 0.2
    height = 0.2

    # 写入标签文件
    label_path = f"{split_dir}/labels/{index}.txt"
    with open(label_path, 'w') as f:
        f.write(f"{label} {x_center} {y_center} {width} {height}")


def create_data_yaml(data_dir):
    """创建 YOLO 数据配置文件"""
    yaml_path = os.path.join(data_dir, 'mnist.yaml')
    with open(yaml_path, 'w') as f:
        f.write(f"""
train: {os.path.abspath(f'{data_dir}/train/images')}
val: {os.path.abspath(f'{data_dir}/val/images')}
test: {os.path.abspath(f'{data_dir}/test/images')}

nc: 10
names: ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
""")
    return yaml_path


# 训练 YOLOv8 模型
def train_yolov8_model(data_yaml, model_size='n', epochs=10, imgsz=416):
    """训练 YOLOv8 模型识别手写数字"""
    # 选择模型大小 (n, s, m, l, x)
    model_name = f'yolov8{model_size}.pt'

    # 加载预训练模型
    model = YOLO(model_name)

    # 训练模型
    model.train(
        # augment=True,  # 启用YOLO内部数据增强
        # scale=0.5,  # 缩放范围
        data=data_yaml,
        epochs=epochs,
        imgsz=imgsz,
        batch=64,
        name=f'yolov8{model_size}_mnist'
    )

    # 评估模型
    metrics = model.val()

    return model, metrics


def load_data():
    convert_mnist_to_yolo()


def test_model(model, image_path):
    """测试训练好的模型"""
    # 加载图像
    # img = cv2.imread(image_path)
    img = Image.open(image_path).resize((800, 800))

    # 进行预测
    results = model(img, imgsz=800)

    # 显示结果
    annotated_img = results[0].plot()

    # 显示图像
    plt.figure(figsize=(10, 10))
    plt.imshow(cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB))
    plt.axis('off')
    plt.show()

    # 打印预测结果
    for box in results[0].boxes:
        class_id = int(box.cls)
        conf = float(box.conf)
        print(f"预测: {class_id}, 置信度: {conf:.2f}")


def train():
    # 转换 MNIST 数据集为 YOLO 格式
    data_dir = "./mnist_yolo"
    data_yaml = os.path.join(data_dir, 'mnist.yaml')

    # # 训练模型 (使用 'n' 轻量版，可改为 's', 'm', 'l', 'x' 获取更高精度)
    train_yolov8_model(data_yaml, model_size='n', epochs=10)

    print("训练结束")


def test():
    model = YOLO('runs/detect/yolov8n_mnist/weights/best.pt')

    # img = Image.open("number/0-9.jpg").resize((416, 416))

    # results = model.predict(source=img, imgsz=416, show=True, save=True)

    # print(results)

    test_model(model, "number/777.jpg")

    time.sleep(5)


if __name__ == "__main__":
   load_data()
   train()



